import datetime
import gradio as gr
import pymongo as pm
from util_mongo import enqueue, get_history, delete_many_by_user
from common import MONGODB_NAME, VALUE, KEY, IO_PREFIX
from util import uuid
import time
from http.cookies import SimpleCookie
from redis_interface import text_gen
import redis

# 连接本机6379端口上的redis的第0个数据库
rdb = redis.Redis('127.0.0.1', 6379, 0, password='lgdz4qEdt/ezElyQnXFYXB80iM3OxEbAWRjMFPcIXH5ni6eQ8QOlfp7G7gvV1svPu2Bv7v')

# 连接Mongodb
mongo = pm.MongoClient(
    '127.0.0.1', 27017, serverSelectionTimeoutMS=3000,
    username='root', password='p1983mdA1-ei',
)
mdb = mongo[MONGODB_NAME]


def chat_once(xinput, username):
    """输入文本，用模型返回输出。"""    
    if username is None:
        # 为了健壮性
        return '', []
    
    # 得到输出
    xoutput = text_gen(xinput, username)
        
    # 输入存入mongodb
    ts = time.time_ns()
    enqueue(mdb, 'dialog_in', username, ts, {
        VALUE: xinput,
        'io': 'i',
    })
    
    # 输出存入mongodb
    ts = time.time_ns()
    enqueue(mdb, 'dialog_out', username, ts, {
        VALUE: xoutput,
        'io': 'o',
    })

    # 获取聊天历史
    xlog = get_history(mdb, username, more_info=True)

    # 返回空字符串清空输入文本框，返回聊天历史
    return '', xlog


def embed_user_id_in_state_comp(request: gr.Request):
    """
    页面加载的事件handler

    从HTTP request中获取cookie来确定用户名，如果没有则生成一个
    :param request: HTTP request data
    :return: 将用户名发送给浏览器中存放用户名的组件
    """
    xuuid = uuid()

    #从cookie中获取username
    username = None
    try:
        cookie_str = request.headers.cookie
        print('cookie_str:', cookie_str)
        cookie = SimpleCookie()
        cookie.load(cookie_str)
        username = cookie.get('username', None)
        print('user:', username)
        username = username.value
    except Exception:
        try:
            username = request.cookies['username']
        except Exception:
            pass

    # 获取不到，新建一个username
    if username is None:
        print('new user!')
        username = f'u{xuuid}'

    # print username
    print('user:', username)

    xlog = get_history(mdb, username, more_info=True)

    return username, username, username, xlog


# 用户名写入cookie的JS
js_username2cookie = """
    async () => {{
            setTimeout(() => {{
                var xnew = document.querySelector('#username.prose').innerHTML;
                console.log('username', xnew)
                if ('' == xnew) {
                    return;
                }
                document.cookie = 'username=' + xnew + '; max-age=' + 3600*24*30*6;
            }}, 1);
        }}
"""


def do_the_init():
    """
        点击”开始聊天“的事件的handler
        :return: 隐藏”开始聊天“，显示其他组件。
    """
    return (
        gr.update(visible=False),
        *([gr.update(visible=True)] * 3),
    )
    
    
def show_text_out_cache(username):
    if username is None:
        # 为了健壮性
        return ''

    cache = rdb.hget('username2dynamic', username)
    if cache is None:
        return ''
    cache = cache.decode('utf8')
    return cache


def clear_history(username):
    if username is None:
        # 为了健壮性
        return []
    delete_many_by_user(mdb, 'dialog_in', username)
    delete_many_by_user(mdb, 'dialog_out', username)
    return []

# 构建Gradio界面
with gr.Blocks(analytics_enabled=False) as demo:

    # 用户名相关控件
    cmp_username_state = gr.State()
    cmp_username_html = gr.HTML(visible=False, elem_id='username')
    cmp_username_display = gr.Textbox(interactive=False, label='用户名')

    # 启动按钮
    cmp_start_btn = gr.Button('开始聊天')

    # 聊天历史
    cmp_chatbot = gr.Chatbot().style(height=350)
    
    # 动态响应输出控件
    cmp_text_out_cache = gr.Textbox(interactive=False, label='动态响应', visible=False, )
    
    # 输入文本框
    cmp_input = gr.Textbox(interactive=True, label='输入', visible=False, )

    # 输入按钮
    cmp_input_button = gr.Button('发送', visible=False, )

    cmp_clear_button = gr.Button('清理历史')

    cmp_clear_button.click(clear_history, cmp_username_state, cmp_chatbot, queue=False)

    # 输入按钮点击事件
    cmp_input_button.click(chat_once, [cmp_input, cmp_username_state], [
        cmp_input,
        cmp_chatbot,
    ], queue=False)

    # 加载时获取或设置cookie，并展示聊天历史
    demo.load(embed_user_id_in_state_comp, None, [
        cmp_username_state,
        cmp_username_display,
        cmp_username_html,
        cmp_chatbot,
    ], queue=False)
    # cmp_username_html发生变化时，触发JS，把生成的username放入页面cookie
    cmp_username_html.change(None, None, None, _js=js_username2cookie, queue=False)
    
    # 启动
    cmp_start_btn.click(do_the_init, None, [
        cmp_start_btn,
        cmp_text_out_cache,
        cmp_input,
        cmp_input_button,
    ], queue=False)
    
    # 输出动态响应
    cmp_start_btn.click(show_text_out_cache, cmp_username_state, cmp_text_out_cache, every=0.1, )

# 使用queue（目前这个例子太简单，没有用到queue，后面会用到。）
demo.queue(
    concurrency_count=10,
    status_update_rate='auto',
    # status_update_rate=0.02,
)

# 启动Gradio
# demo.launch(server_name='0.0.0.0', server_port=6006, share=True, debug=True)  # 带内网穿透
demo.launch(server_name='0.0.0.0', server_port=6006, debug=True)
